package com.lan.ai.mutimodelEmbedding;

import cn.hutool.json.JSONUtil;
import org.springframework.ai.document.ContentFormatter;
import org.springframework.ai.document.Document;
import org.springframework.ai.document.MetadataMode;
import org.springframework.ai.document.id.IdGenerator;
import org.springframework.ai.document.id.RandomIdGenerator;
import org.springframework.ai.model.Media;
import org.springframework.util.Assert;
import org.springframework.util.MimeType;

import java.util.Base64;
import java.util.Collection;
import java.util.List;
import java.util.Map;

/**
 * @author lanyanhua
 * @date 2025/1/21 14:43
 * @description
 */
public class MutimodeDocument extends Document {
    public MutimodeDocument(String content) {
        super(content);
    }

    public MutimodeDocument(String content, Map<String, Object> metadata) {
        this(content, metadata, new RandomIdGenerator());
    }

    public MutimodeDocument(String content, Collection<Media> media, Map<String, Object> metadata) {
        this(new RandomIdGenerator().generateId(content, metadata), content, media, metadata);
    }

    public MutimodeDocument(String content, Map<String, Object> metadata, IdGenerator idGenerator) {
        this(idGenerator.generateId(content, metadata), content, metadata);
    }

    public MutimodeDocument(String id, String content, Map<String, Object> metadata) {
        this(id, content, List.of(), metadata);

    }

    public MutimodeDocument(String id, String content, Collection<Media> media, Map<String, Object> metadata) {
        super(id, content, media, metadata);
        Assert.isTrue(media.isEmpty() || media.size() == 1
                || media.stream().allMatch(m -> m.getMimeType().getType().equals("image")), "media 只能有一个且为图片");
    }


    @Override
    public String getContent() {
        if (getMedia().isEmpty()) {
            return JSONUtil.toJsonStr(List.of(MutimodelEmbeddingRequest.text(super.getContent())));
        }
        return JSONUtil.toJsonStr(getMedia().stream().filter(media -> media.getMimeType().getType().equals("image"))
                .map(media -> MutimodelEmbeddingRequest.contents(null,
//                .map(media -> MutimodelEmbeddingRequest.contents(super.getContent(), 不能携带文本内容 模型图片、文本好像只能二取一
                        fromMediaData(media.getMimeType(), media.getData())
                        , null
                ))
                .toList());
    }

    public String sourceContent() {
        //返回 原始内容+图片信息
        return super.getContent();
//        return JSONUtil.toJsonStr(getMedia().stream().filter(media -> media.getMimeType().getType().equals("image"))
//                .map(media -> MutimodelEmbeddingRequest.contents(,
//                        JSONUtil.toJsonStr(media)
//                        , null
//                ))
//                .toList());
    }


    public String fromMediaData(MimeType mimeType, Object mediaContentData) {
        if (mediaContentData instanceof byte[] bytes) {
            // Assume the bytes are an image. So, convert the bytes to a base64 encoded
            // following the prefix pattern.
            return String.format("data:%s;base64,%s", mimeType.toString(), Base64.getEncoder().encodeToString(bytes));
        } else if (mediaContentData instanceof String text) {
            // Assume the text is a URLs or a base64 encoded image prefixed by the user.
            return text;
        } else {
            throw new IllegalArgumentException(
                    "Unsupported media data type: " + mediaContentData.getClass().getSimpleName());
        }
    }

//    @JsonIgnore
//    public String getFormattedContent() {
//        return getContent();
//    }
//
//    public String getFormattedContent(MetadataMode metadataMode) {
////        Assert.notNull(metadataMode, "Metadata mode must not be null");
//        return getContent();
//    }
//

    /**
     * Helper content extractor that uses and external {@link ContentFormatter}.
     */
    public String getFormattedContent(ContentFormatter formatter, MetadataMode metadataMode) {
//        Assert.notNull(formatter, "formatter must not be null");
//        Assert.notNull(metadataMode, "Metadata mode must not be null");
        return super.getContent();
    }


}
